
import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.asn1.x500.X500Name;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
import org.bouncycastle.cert.X509v3CertificateBuilder;
import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumKeyPairGenerator;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumParameters;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumPrivateKeyParameters;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumPublicKeyParameters;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumSigner;
import org.bouncycastle.pqc.jcajce.provider.BouncyCastlePQCProvider;
import org.bouncycastle.util.io.pem.PemObject;
import org.bouncycastle.util.io.pem.PemWriter;
import java.io.FileOutputStream;
import java.io.FileWriter;
import java.io.IOException;
import java.math.BigInteger;
import java.security.cert.X509Certificate;
import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.Date;
import javax.crypto.*;
import javax.crypto.spec.*;
import java.security.*;

import org.bouncycastle.pqc.crypto.crystals.dilithium.*;
import org.bouncycastle.pqc.crypto.util.PrivateKeyInfoFactory;
import org.bouncycastle.pqc.crypto.util.SubjectPublicKeyInfoFactory;
import org.bouncycastle.asn1.pkcs.PrivateKeyInfo;
import org.bouncycastle.asn1.pkcs.*;




public class Main {
    //注册加密提供者普通密码学和后量子密码学
    static {
        Security.addProvider(new BouncyCastleProvider());
        Security.addProvider(new BouncyCastlePQCProvider());
    }

    public static void main(String[] args) throws Exception {


        try {
            System.out.println("使用底层Dilithium API生成证书...");
            // 测试ML-DSA-44，ML-DSA-65和ML-DSA-87
            System.out.println("=== ML-DSA-44 生成X509证书示例 ===");
            generateCertificateWithLowLevelAPI(2);

            System.out.println("=== ML-DSA-65 生成X509证书示例 ===");
            generateCertificateWithLowLevelAPI(3);

            System.out.println("=== ML-DSA-87 生成X509证书示例 ===");
            generateCertificateWithLowLevelAPI(5);

        } catch (Exception e) {
            System.out.println("证书生成失败: " + e.getMessage());
            e.printStackTrace();
        }
    }

    private static void generateCertificateWithLowLevelAPI(Integer KeyType) throws Exception {
        SecureRandom random = new SecureRandom();

        // 默认使用Dilithium2参数
        DilithiumParameters params = DilithiumParameters.dilithium5;
        if(KeyType==2){
            params = DilithiumParameters.dilithium2;
        }
        if(KeyType==3){
            params = DilithiumParameters.dilithium3;
        }
        if(KeyType==5){
            params = DilithiumParameters.dilithium5;
        }
        DilithiumKeyPairGenerator keyPairGen = new DilithiumKeyPairGenerator();
        keyPairGen.init(new org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumKeyGenerationParameters(random, params));

        var keyPair = keyPairGen.generateKeyPair();
        DilithiumPublicKeyParameters pubKey = (DilithiumPublicKeyParameters) keyPair.getPublic();
        DilithiumPrivateKeyParameters privKey = (DilithiumPrivateKeyParameters) keyPair.getPrivate();


        System.out.println("Dilithium"+ KeyType +" 密钥对生成成功");

        // 设置正确的OID和参数（包含NULL）
        AlgorithmIdentifier sigAlgId = getAlgorithmIdentifier(params);
        System.out.println("使用算法OID: " + sigAlgId.getAlgorithm());

        PrivateKeyInfo prikeypkcs8 = new PrivateKeyInfo(sigAlgId,privKey.getEncoded());
        SubjectPublicKeyInfo pubkeypkcs8 = new SubjectPublicKeyInfo(sigAlgId,pubKey.getEncoded());

        // 构建证书
        X500Name subject = new X500Name("CN=Test Dilithium"+ KeyType +" Certificate, O=Demo, C=CN");
        BigInteger serial = BigInteger.valueOf(System.currentTimeMillis());
        Instant now = Instant.now();
        Date notBefore = Date.from(now);
        Date notAfter = Date.from(now.plus(365, ChronoUnit.DAYS));

        // 公钥信息 - 使用正确的ASN.1编码
        SubjectPublicKeyInfo pubInfo = createSubjectPublicKeyInfo(pubKey, params);

        X509v3CertificateBuilder certGen = new X509v3CertificateBuilder(
                subject, serial, notBefore, notAfter, subject, pubInfo);

        // 自定义内容签名器
        ContentSigner signer = new ContentSigner() {
            private final java.io.ByteArrayOutputStream buffer = new java.io.ByteArrayOutputStream();

            @Override
            public AlgorithmIdentifier getAlgorithmIdentifier() {
                return sigAlgId;
            }

            @Override
            public java.io.OutputStream getOutputStream() {
                return buffer;
            }

            @Override
            public byte[] getSignature() {
                try {
                    byte[] dataToSign = buffer.toByteArray();
                    DilithiumSigner dilithiumSigner = new DilithiumSigner();
                    dilithiumSigner.init(true, privKey);
                    byte[] signature = dilithiumSigner.generateSignature(dataToSign);

                    // Dilithium签名通常已经是ASN.1编码格式，但确保正确
                    return signature;
                } catch (Exception e) {
                    throw new RuntimeException("签名失败: " + e.getMessage(), e);
                }
            }
        };

        System.out.println("开始构建证书...");
        org.bouncycastle.cert.X509CertificateHolder certHolder = certGen.build(signer);
        System.out.println("证书构建完成");

        // 转换为JCA证书
        X509Certificate cert = new JcaX509CertificateConverter()
                .setProvider("BC")
                .getCertificate(certHolder);


        System.out.println("证书生成成功");


        // 输出文件
        writeCertificateToPEM(cert, "./out/dilithium"+ KeyType +"_cert.crt");
        writePrivateKeyToPEM(prikeypkcs8.getEncoded(), "./out/dilithium"+ KeyType +"_private_key.pem");
        writePublicKeyToPEM(pubkeypkcs8.getEncoded(), "./out/dilithium"+ KeyType +"_public_key.pem");
        writeCertificateToDER(cert, "./out/dilithium"+ KeyType +"_cert.der");

        System.out.println("=== 证书信息 ===");
        System.out.println("主题: " + cert.getSubjectX500Principal());
        System.out.println("颁发者: " + cert.getIssuerX500Principal());
        System.out.println("签名算法: " + cert.getSigAlgName());
        System.out.println("序列号: " + cert.getSerialNumber());
        System.out.println("有效期: " + cert.getNotBefore() + " 至 " + cert.getNotAfter());
        System.out.println("所有文件已成功生成！");
    }

    private static AlgorithmIdentifier getAlgorithmIdentifier(DilithiumParameters params) {
        // 返回包含NULL参数的算法标识符
        ASN1ObjectIdentifier oid = getAlgorithmOID(params);
        return new AlgorithmIdentifier(oid); // 添加NULL参数, DERNull.INSTANCE
    }

    private static ASN1ObjectIdentifier getAlgorithmOID(DilithiumParameters params) {
        // 根据Dilithium参数返回正确的OID
        if (params == DilithiumParameters.dilithium2) {
            return new ASN1ObjectIdentifier("2.16.840.1.101.3.4.3.17");
        } else if (params == DilithiumParameters.dilithium3) {
            return new ASN1ObjectIdentifier("2.16.840.1.101.3.4.3.18");
        } else if (params == DilithiumParameters.dilithium5) {
            return new ASN1ObjectIdentifier("2.16.840.1.101.3.4.3.19");
        } else {
            return new ASN1ObjectIdentifier("2.16.840.1.101.3.4.3.19");
        }
    }

    private static SubjectPublicKeyInfo createSubjectPublicKeyInfo(
            DilithiumPublicKeyParameters pubKey, DilithiumParameters params) {
        try {
            AlgorithmIdentifier algorithmId = new AlgorithmIdentifier(
                    getAlgorithmOID(params)); // 公钥算法标识符也包含NULL ,DERNull.INSTANCE

            return new SubjectPublicKeyInfo(algorithmId, pubKey.getEncoded());
        } catch (Exception e) {
            throw new RuntimeException("创建公钥信息失败", e);
        }
    }

    // PEM输出方法保持不变
    private static void writeCertificateToPEM(X509Certificate cert, String filename) throws Exception {
        try (PemWriter pemWriter = new PemWriter(new FileWriter(filename))) {
            PemObject pemObject = new PemObject("CERTIFICATE", cert.getEncoded());
            pemWriter.writeObject(pemObject);
        }
        System.out.println("证书已保存: " + filename);
    }

    private static void writePrivateKeyToPEM(byte[] privateKeyBytes, String filename) throws IOException {
        try (PemWriter pemWriter = new PemWriter(new FileWriter(filename))) {
            PemObject pemObject = new PemObject("PRIVATE KEY", privateKeyBytes);
            pemWriter.writeObject(pemObject);
        }
        System.out.println("私钥已保存: " + filename);
    }

    private static void writePublicKeyToPEM(byte[] publicKeyBytes, String filename) throws IOException {
        try (PemWriter pemWriter = new PemWriter(new FileWriter(filename))) {
            PemObject pemObject = new PemObject("PUBLIC KEY", publicKeyBytes);
            pemWriter.writeObject(pemObject);
        }
        System.out.println("公钥已保存: " + filename);
    }

    private static void writeCertificateToDER(X509Certificate cert, String filename) throws Exception {
        try (FileOutputStream fos = new FileOutputStream(filename)) {
            fos.write(cert.getEncoded());
        }
        System.out.println("DER证书已保存: " + filename);
    }
}
